package com.starhub.integration.sdk.baidu.model;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.output.Response;
import lombok.Builder;
import org.springframework.beans.factory.annotation.Configurable;

import com.starhub.integration.sdk.baidu.core.BaiduClient;
import com.starhub.integration.sdk.baidu.core.BaiduClientConnector;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 百度向量化模型
 */
@Builder
@Configurable
public class BaiduEmbeddings implements EmbeddingModel {
    
    private static volatile BaiduClientConnector clientConnector;
    private static final Object lock = new Object();
    
    private BaiduClientConnector getClientConnector() {
        if (clientConnector == null) {
            synchronized (lock) {
                if (clientConnector == null) {
                    clientConnector = new BaiduClientConnector();
                }
            }
        }
        return clientConnector;
    }

    private final String mark;
    private final String apiKey;
    private final String secretKey;
    
    /**
     * 构造函数
     * @param mark 模型标识
     * @param apiKey 百度API Key
     * @param secretKey 百度Secret Key
     */
    public BaiduEmbeddings(String mark, String apiKey, String secretKey) {
        this.mark = mark;
        this.apiKey = apiKey;
        this.secretKey = secretKey;
    }
    
    /**
     * 向量化
     * @param textSegment 文本段
     * @return 向量
     */
    @Override
    public Response<Embedding> embed(TextSegment textSegment) {
        try {
            BaiduClient baiduClient = getClientConnector().getBaiduClient(mark, apiKey, secretKey);
            float[] vector = baiduClient.getEmbedding(textSegment.text());
            return Response.from(new Embedding(vector));
        } catch (IOException e) {
            throw new RuntimeException("文本转向量失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 向量化
     * @param text 文本
     * @return 向量
     */
    @Override
    public Response<Embedding> embed(String text) {
        try {
            BaiduClient baiduClient = getClientConnector().getBaiduClient(mark, apiKey, secretKey);
            float[] vector = baiduClient.getEmbedding(text);
            return Response.from(new Embedding(vector));
        } catch (IOException e) {
            throw new RuntimeException("文本转向量失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 向量化
     * @param textSegments 文本段列表
     * @return 向量列表
     */
    @Override
    public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
        List<Embedding> embeddings = new ArrayList<>();
        try {
            BaiduClient baiduClient = getClientConnector().getBaiduClient(mark, apiKey, secretKey);
            for (TextSegment segment : textSegments) {
                float[] vector = baiduClient.getEmbedding(segment.text());
                embeddings.add(new Embedding(vector));
            }
            return Response.from(embeddings);
        } catch (IOException e) {
            throw new RuntimeException("文本转向量失败: " + e.getMessage(), e);
        }
    }

    public String getMark() {
        return mark;
    }

    public String getApiKey() {
        return apiKey;
    }

    public String getSecretKey() {
        return secretKey;
    }

    public static BaiduEmbeddingsBuilder builder() {
        return new BaiduEmbeddingsBuilder();
    }

    public static class BaiduEmbeddingsBuilder {
        private String mark;
        private String apiKey;
        private String secretKey;

        public BaiduEmbeddingsBuilder mark(String mark) {
            this.mark = mark;
            return this;
        }

        public BaiduEmbeddingsBuilder apiKey(String apiKey) {
            this.apiKey = apiKey;
            return this;
        }

        public BaiduEmbeddingsBuilder secretKey(String secretKey) {
            this.secretKey = secretKey;
            return this;
        }

        public BaiduEmbeddings build() {
            return new BaiduEmbeddings(mark, apiKey, secretKey);
        }
    }
}